"""
--------------------------------------------------------------------------------
功能:节点管理程序部分，该程序可调用
    1.serial_section:和下位机通信的串口管理部分
    2.servo_section:和下位机通信的舵机协议部分
    3.user_arm:机械臂逆运动学部分
    实现单个程序管理三个节点, 实现控制机械臂运动

日期:2025-06-03
作者:Unirobot
QQ交流群:811348489
寄语：代码中能遇见你，真好，有你，真好.
署名：有你同创智能机器人科技(北京)有限公司.
--------------------------------------------------------------------------------
"""
import time
from src.arm_control import serial_section, servo_section, user_arm

class RobotManager:
    def __init__(self):
        #创建并打开和下位机通信串口对象
        self.serial = serial_section.SerialPort(port="/dev/ttyS1", baudrate=115200)
        self.serial.open()
        time.sleep(0.2)
        #创建舵机协议对象
        self.servos = servo_section.ServoProtocol()
        time.sleep(0.2)
        #创建机械臂对象
        self.arm = user_arm.FourDofArm()
        time.sleep(0.2)
        #定义四自由度机械臂参数
        self.angle = [0.0, 0.0, 0.0, 0.0]
        self.start_ang = [50.0, 0.0, 80.0]  #机械臂逆运动学初始化参数
        self.save_ang = [50.0, 0.0, 80.0]   #实时保存四自由度机械臂逆运动学角度
        for i in range(3):
            self.save_ang[i] = self.start_ang[i]

        #设置机械臂角度最大值
        self.arm_limit_x_min = 50
        self.arm_limit_x_max = 180
        self.arm_limit_y_min = -80
        self.arm_limit_y_max = 80
        self.arm_limit_z_min = -50
        self.arm_limit_z_max = 140

        self.gimbal_limit_hor_min = 30
        self.gimbal_limit_hor_max = 200
        self.gimbal_limit_ver_min = 30
        self.gimbal_limit_ver_max = 200

        # # 机械臂执行初始动作
        # self.release_sth() #打开爪子
        # time.sleep(0.6)
        # #机械臂执行初始化动作
        # self.arm_init_action()
        # time.sleep(0.6)

        self.arm_move_x = self.start_ang[0]
        self.arm_move_y = self.start_ang[1]
        self.arm_move_z = self.start_ang[2]
        self.gimbal_ang = [140, 140]
        # self.gimbal_move(140, 140, 10)
        # time.sleep(1)
        # self.gimbal_move(90, 90, 10)
        # time.sleep(1)
        # self.gimbal_move(140, 140, 10)
        # time.sleep(1)

    """
    终端打印出带颜色的字体(主要是警告，或者是错误信息)
    """
    def color_text(self, text, color=None, bg_color=None, bold=False, underline=False):
        codes = []
        if color:
            color_codes = {
                "black": 30, "red": 31, "green": 32, "yellow": 33,
                "blue": 34, "magenta": 35, "cyan": 36, "white": 37,
            }
            codes.append(color_codes.get(color, 39))
        if bg_color:
            bg_codes = {
                "black": 40, "red": 41, "green": 42, "yellow": 43,
                "blue": 44, "magenta": 45, "cyan": 46, "white": 47,
            }
            codes.append(bg_codes.get(bg_color, 49))
        if bold:
            codes.append(1)
        if underline:
            codes.append(4)
        
        if codes:
            return f"\033[{';'.join(map(str, codes))}m{text}\033[0m"
        else:
            return text

    def arm_limit_x(self, input_value_x):
        if input_value_x <= self.arm_limit_x_min:
            # print(self.color_text("机械臂x轴到达最小值", color="red", bold=True))
            input_value_x = self.arm_limit_x_min
        if input_value_x >= self.arm_limit_x_max:
            # print(self.color_text("机械臂x轴到达最小值", color="red", bold=True))
            input_value_x = self.arm_limit_x_max
        return input_value_x

    def arm_limit_y(self, input_value_y):
        if input_value_y <= self.arm_limit_y_min:
            # print(self.color_text("机械臂y轴到达最小值", color="red", bold=True))
            input_value_y = self.arm_limit_y_min
        if input_value_y >= self.arm_limit_y_max:
            # print(self.color_text("机械臂y轴到达最小值", color="red", bold=True))
            input_value_y = self.arm_limit_y_max
        return input_value_y

    def arm_limit_z(self, input_value_z):
        if input_value_z <= self.arm_limit_z_min:
            # print(self.color_text("机械臂z轴到达最小值", color="red", bold=True))
            input_value_z = self.arm_limit_z_min
        if input_value_z >= self.arm_limit_z_max:
            # print(self.color_text("机械臂z轴到达最小值", color="red", bold=True))
            input_value_z = self.arm_limit_z_max
        return input_value_z

    def gimbal_limit_hor(self, input_value_hor):
        if input_value_hor <= self.gimbal_limit_hor_min:
            # print(self.color_text("机械臂x轴到达最小值", color="red", bold=True))
            input_value_hor = self.gimbal_limit_hor_min
        if input_value_hor >= self.gimbal_limit_hor_max:
            # print(self.color_text("机械臂x轴到达最小值", color="red", bold=True))
            input_value_hor = self.gimbal_limit_hor_max
        return input_value_hor

    def gimbal_limit_ver(self, input_value_ver):
        if input_value_ver <= self.gimbal_limit_ver_min:
            # print(self.color_text("机械臂x轴到达最小值", color="red", bold=True))
            input_value_ver = self.gimbal_limit_ver_min
        if input_value_ver >= self.gimbal_limit_ver_max:
            # print(self.color_text("机械臂x轴到达最小值", color="red", bold=True))
            input_value_ver = self.gimbal_limit_ver_max
        return input_value_ver

    def three_axis_data_limit(self, input_value_x, input_value_y, input_value_z):
        if input_value_x <= self.arm_limit_x_min:
            input_value_x = self.arm_limit_x_min
        if input_value_x >= self.arm_limit_x_max:
            input_value_x = self.arm_limit_x_max

        if input_value_y <= self.arm_limit_y_min:
            input_value_y = self.arm_limit_y_min
        if input_value_y >= self.arm_limit_y_max:
            input_value_y = self.arm_limit_y_max

        if input_value_z <= self.arm_limit_z_min:
            input_value_z = self.arm_limit_z_min
        if input_value_z >= self.arm_limit_z_max:
            input_value_z = self.arm_limit_z_max
        return input_value_x, input_value_y, input_value_z
    """
    将x从区间(a,b)线性映射到区间(c,d)
    参数:
        x: 要映射的值
        a, b: 原区间
        c, d: 目标区间
    返回:
        映射后的值
    """
    def linear_mapping(self, x, a, b, c, d):
        return (x - a) * (d - c) / (b - a) + c

        

    """
    通过线性方程将物体坐标转换为机械臂抓取的逆运动学坐标
    """
    def calculate_arm_pos(self, x, y):
        position_y = 0.3638 * x - 15.049
        position_x = 0.821 * y + 114.2
        position_x = round(position_x, 2)
        position_y = round(position_y, 2)
        # print("position_x:", position_x, "position_y:", position_y)

        position_x = self.arm_limit_x(position_x)
        position_y = self.arm_limit_y(position_y)
        return position_x, position_y

    """
    机械臂手爪闭合
    """
    def catch_sth(self):
        data = self.servos.servo_control(
            servo_num=1,
            servo_id=[6],
            servo_angle=[110]
        )
        self.serial.serial_send(data)
        time.sleep(0.05)

    """
    机械臂手爪打开
    """
    def release_sth(self):
        data = self.servos.servo_control(
            servo_num=1,
            servo_id=[6],
            servo_angle=[1.01]
        )
        self.serial.serial_send(data)
        time.sleep(0.05)

    """
    添加云台控制
    """
    def gimbal_move(self, Horizontal_angle, Vertical_angle, input_point):

        num_points = input_point
        for i in range(num_points + 1):
            t = i / num_points
            x = self.gimbal_ang[0] + t * (Horizontal_angle - self.gimbal_ang[0])
            y = self.gimbal_ang[1] + t * (Vertical_angle  -  self.gimbal_ang[1])

            x = self.gimbal_limit_hor(x)
            y = self.gimbal_limit_ver(y)

            data = self.servos.servo_control(
                servo_num=2,
                servo_id=[3, 4],
                servo_angle=[x, y]
            )
            self.serial.serial_send(data)
            time.sleep(0.03)

        self.gimbal_ang[0] = Horizontal_angle
        self.gimbal_ang[1] = Vertical_angle    




    """
    机械臂移动函数
    该函数只需要传入机械臂逆运动学的目标位置即可.
    意思就是: 该函数会自动记录机械臂上一次机械臂的目标位置, 
            传入目标位置后, 机械臂会自动从上一次位置运动到新传入的目标位置
    """
    def arm_move(self, target_x, target_y, target_z, input_point):
        current_angle = [0.0, 0.0, 0.0]
        delta_x = self.save_ang[0] - target_x
        delta_y = self.save_ang[1] - target_y
        num_points = 0
        #防止数据出错
        # try:
        #     num_points = int((math.fabs(delta_x) + math.fabs(delta_y)) / 9.0)
        # except:
        #     num_points = 1
        # if num_points <= 3:
        #     num_points = 3
        # if num_points >= 6:
        #     num_points = 6
        num_points = input_point
        for i in range(num_points + 1):
            t = i / num_points
            x = self.save_ang[0] + t * (target_x - self.save_ang[0])
            y = self.save_ang[1] + t * (target_y - self.save_ang[1])
            z = self.save_ang[2] + t * (target_z - self.save_ang[2])
            x = self.arm_limit_x(x)
            y = self.arm_limit_y(y)
            z = self.arm_limit_z(z)

            current_angle[0] = x
            current_angle[1] = y
            current_angle[2] = z

            self.angle[0],self.angle[1],self.angle[2],self.angle[3] =  self.arm.get_coordinate(x, y, z)
            data = self.servos.servo_control(
                servo_num=4,
                servo_id=[44, 45, 46, 7],
                servo_angle=[self.angle[3], self.angle[2], self.angle[1], self.angle[0]]
            )
            self.serial.serial_send(data)
            time.sleep(0.03)
        self.save_ang = [current_angle[0], current_angle[1], current_angle[2]]

    """
    机械臂初始化动作
    机械臂从上一次的角度转动到初始化动作
    """
    def arm_init_action(self):
        self.release_sth() #打开爪子
        # time.sleep(0.6)
        self.arm_move(self.start_ang[0], self.start_ang[1], self.start_ang[2], 2)
        # time.sleep(0.6)
